"""
PE文件头特征提取模块
功能：从Windows可执行文件中提取关键头信息
"""

import pefile

def extract_pe_features(file_path):
    """提取PE文件特征
    参数：
        file_path: PE文件路径
    返回：
        dict 包含提取的特征字典，解析失败返回None
    """
    features = {}
    try:
        pe = pefile.PE(file_path)
        
        # 基础头信息
        features['Machine'] = pe.FILE_HEADER.Machine  # 机器类型
        features['SizeOfCode'] = pe.OPTIONAL_HEADER.SizeOfCode  # 代码段大小
        features['DllCharacteristics'] = pe.OPTIONAL_HEADER.DllCharacteristics  # DLL特征
        
        # 导入表分析
        if hasattr(pe, 'DIRECTORY_ENTRY_IMPORT'):
            features['ImportedDLLs'] = len(pe.DIRECTORY_ENTRY_IMPORT)  # 导入DLL数量
            api_calls = []
            for entry in pe.DIRECTORY_ENTRY_IMPORT:
                for imp in entry.imports:
                    if imp.name:
                        api_calls.append(imp.name.decode())
            features['UniqueAPIs'] = len(set(api_calls))  # 唯一API数量
        else:
            features['ImportedDLLs'] = 0
            features['UniqueAPIs'] = 0
            
    except Exception as e:
        print(f"文件解析失败: {file_path} - 错误: {str(e)}")
        return None
    
    return features 